# test/simple_memory_test.py
import os
import shutil
import json
import uuid
import time
import sqlite3
from unittest.mock import Mock, patch # Keep patch for threading mock

# Import real implementations
from root_agent.llm import LLMInterface
from root_agent.utils.llm_wrapper import GMemoryLLMWrapper
from memory_systems.GMemory.mas.memory.mas_memory.GMemory import GMemory
from memory_systems.GMemory.mas.utils import EmbeddingFunc, get_model_type
from memory_systems.vector_db import ZillizVectorDB as VectorDB # ZillizVectorDB remains mocked for now

# Adjust the path to import modules from the project root
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from root_agent.utils.memory_management import MemoryManager
from memory_systems.experience_metadata_db import ExperienceMetadataDB

# Define a temporary database path for testing
# Use an absolute path relative to the project root
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
TEST_DB_DIR = os.path.join(PROJECT_ROOT, "memory_systems")
TEST_DB_PATH = os.path.join(TEST_DB_DIR, "test_simple_memory.db")
TEST_LLM_LOG_DIR = os.path.join(PROJECT_ROOT, "temp", "test_simple_llm_logs")
TEST_GMEMORY_WORKING_DIR = os.path.join(PROJECT_ROOT, ".db", "test_gmemory")


def setup_test_environment():
    """Sets up a clean test environment."""
    os.makedirs(TEST_DB_DIR, exist_ok=True) # Ensure the directory for the DB exists
    if os.path.exists(TEST_DB_PATH):
        os.remove(TEST_DB_PATH)
    if os.path.exists(TEST_LLM_LOG_DIR):
        shutil.rmtree(TEST_LLM_LOG_DIR)
    os.makedirs(TEST_LLM_LOG_DIR, exist_ok=True)
    if os.path.exists(TEST_GMEMORY_WORKING_DIR):
        shutil.rmtree(TEST_GMEMORY_WORKING_DIR)
    os.makedirs(TEST_GMEMORY_WORKING_DIR, exist_ok=True)

def teardown_test_environment():
    """Cleans up the test environment."""
    # if os.path.exists(TEST_DB_PATH):
    #     os.remove(TEST_DB_PATH) # Commented out to keep the DB file for inspection
    if os.path.exists(TEST_LLM_LOG_DIR):
        shutil.rmtree(TEST_LLM_LOG_DIR)
    if os.path.exists(TEST_GMEMORY_WORKING_DIR):
        shutil.rmtree(TEST_GMEMORY_WORKING_DIR)


def run_simple_memory_test():
    setup_test_environment()
    print("--- Starting Simple Memory Test ---")

    # 1. Instantiate real LLMInterface
    real_llm = LLMInterface()
    real_llm.set_log_dir(TEST_LLM_LOG_DIR) # Set log directory for real LLM

    # 2. Instantiate real GMemory
    # GMemory requires an LLM model and an embedding function
    gmemory_llm_wrapper = GMemoryLLMWrapper(real_llm)
    embedding_func = EmbeddingFunc()
    
    # Ensure the GMemory working directory is unique for this test run
    gmemory_working_dir_instance = os.path.join(TEST_GMEMORY_WORKING_DIR, str(uuid.uuid4()))
    os.makedirs(gmemory_working_dir_instance, exist_ok=True)

    real_gmemory = GMemory(
        llm_model=gmemory_llm_wrapper,
        embedding_func=embedding_func,
        namespace="gmemory_test",
        global_config={
            "hop": 1,
            "start_insights_threshold": 5,
            "rounds_per_insights": 5,
            "insights_point_num": 5,
            "working_dir": gmemory_working_dir_instance
        }
    )

    # 3. Mock ZillizVectorDB (semantic_memory) as it requires external setup
    mock_semantic_memory = Mock(spec=VectorDB)
    mock_semantic_memory.upsert_text_zilliz.return_value = None
    # The real EmbeddingFunc is used by ZillizVectorDB, so we pass it to the mock
    mock_semantic_memory.embedding_func = EmbeddingFunc() 
    mock_semantic_memory.collection_name = "test_collection"
    
    # Mock the search_vectors return value for retrieval test
    # We'll dynamically set the ID later to match what's inserted
    mock_semantic_memory.search_vectors.return_value = [] 

    # Create a simplified mock root_agent_instance
    # Only include attributes that MemoryManager actually accesses
    mock_root_agent_instance = Mock()
    mock_root_agent_instance.use_gmemory = True
    mock_root_agent_instance.use_vector_memory = True
    mock_root_agent_instance.llm = real_llm # Use real LLM
    mock_root_agent_instance.semantic_memory = mock_semantic_memory # Use mocked Zilliz
    mock_root_agent_instance.gmemory = real_gmemory # Use real GMemory
    mock_root_agent_instance.llm_usage_metrics = {
        "total_input_tokens": 0,
        "total_output_tokens": 0,
        "total_cost": 0.0,
        "task_llm_metrics": []
    }

    # Instantiate MemoryManager and ExperienceMetadataDB directly
    memory_manager = MemoryManager(mock_root_agent_instance)
    # Override the default db_path to use our test path
    memory_manager.experience_metadata_db = ExperienceMetadataDB(TEST_DB_PATH) 

    # --- Test Case 1: Add a successful memory ---
    print("\n--- Test Case 1: Adding a successful memory ---")
    user_goal_success = "Analyze sales data and generate report"
    plan_success = [{"objective": "get_data", "tool_name": "read_file"}]
    results_success = ["data_read_successfully"]
    total_duration_success = 15.5
    total_cost_success = 0.123

    # Call the asynchronous method directly for testing purposes
    # In real usage, this would be called via memory_manager.add_memory() which spawns a thread
    with patch('threading.Thread.start', side_effect=lambda x: x.run()):
        memory_manager.add_memory(
            user_goal_success, plan_success, results_success, True, 
            "user123", "openid456", total_duration_success, total_cost_success
        )
    time.sleep(0.1) # Give a moment for the async operation to potentially complete if not mocked

    # Verify upsert_text_zilliz was called and get the summary_id
    mock_semantic_memory.upsert_text_zilliz.assert_called_once()
    upsert_args, upsert_kwargs = mock_semantic_memory.upsert_text_zilliz.call_args
    summary_id_success = upsert_kwargs['Id']
    print(f"Generated summary_id for success: {summary_id_success}")

    # Verify metadata in SQLite DB
    metadata_success = memory_manager.experience_metadata_db.get_experience_metadata(summary_id_success)
    print(f"Retrieved metadata for success: {metadata_success}")
    assert metadata_success is not None
    assert metadata_success['success_count'] == 1
    assert metadata_success['failure_count'] == 0
    assert metadata_success['total_references'] == 1
    assert abs(metadata_success['average_cost'] - total_cost_success) < 0.0001
    assert abs(metadata_success['average_duration'] - total_duration_success) < 0.0001
    print("Test Case 1 Passed: Successful memory added and metadata verified.")

    # --- Test Case 2: Add a failed memory ---
    print("\n--- Test Case 2: Adding a failed memory ---")
    user_goal_fail = "Fix bug in authentication module"
    plan_fail = [{"objective": "debug_code", "tool_name": "execute_command"}]
    results_fail = ["error_during_execution"]
    total_duration_fail = 5.0
    total_cost_fail = 0.01

    with patch('threading.Thread.start', side_effect=lambda x: x.run()):
        memory_manager.add_memory(
            user_goal_fail, plan_fail, results_fail, False, 
            "user123", "openid456", total_duration_fail, total_cost_fail
        )
    time.sleep(0.1)

    # Verify upsert_text_zilliz was called again and get the summary_id
    # Note: call_args_list[1] because it's the second call
    upsert_args_fail, upsert_kwargs_fail = mock_semantic_memory.upsert_text_zilliz.call_args_list[1]
    summary_id_fail = upsert_kwargs_fail['Id']
    print(f"Generated summary_id for failure: {summary_id_fail}")

    # Verify metadata in SQLite DB
    metadata_fail = memory_manager.experience_metadata_db.get_experience_metadata(summary_id_fail)
    print(f"Retrieved metadata for failure: {metadata_fail}")
    assert metadata_fail is not None
    assert metadata_fail['success_count'] == 0
    assert metadata_fail['failure_count'] == 1
    assert metadata_fail['total_references'] == 1
    assert abs(metadata_fail['average_cost'] - total_cost_fail) < 0.0001
    assert abs(metadata_fail['average_duration'] - total_duration_fail) < 0.0001
    print("Test Case 2 Passed: Failed memory added and metadata verified.")

    # --- Test Case 3: Retrieve memory with metadata ---
    print("\n--- Test Case 3: Retrieving memory with metadata ---")
    # Configure mock_semantic_memory.search_vectors to return our successful summary
    mock_zilliz_result = Mock()
    mock_zilliz_result.id = summary_id_success
    mock_zilliz_result.entity = {"text": "Mock Summary for Sales Data Analysis.", "type": "experience_summary"}
    mock_semantic_memory.search_vectors.return_value = [[mock_zilliz_result]]

    retrieved_memories = memory_manager.retrieve_memory(user_goal_success)
    print(f"Retrieved memories: {retrieved_memories}")

    assert "semantic_summaries" in retrieved_memories
    assert len(retrieved_memories["semantic_summaries"]) == 1
    
    retrieved_summary_item = retrieved_memories["semantic_summaries"][0]
    assert "summary" in retrieved_summary_item
    assert retrieved_summary_item["summary"] == "Mock Summary for Sales Data Analysis."
    
    assert "metadata" in retrieved_summary_item
    retrieved_metadata = retrieved_summary_item["metadata"]
    assert retrieved_metadata is not None
    assert retrieved_metadata['experience_id'] == summary_id_success
    assert retrieved_metadata['success_count'] == 1
    assert retrieved_metadata['failure_count'] == 0
    assert retrieved_metadata['total_references'] == 1
    assert abs(retrieved_metadata['average_cost'] - total_cost_success) < 0.0001
    assert abs(retrieved_metadata['average_duration'] - total_duration_success) < 0.0001
    print("Test Case 3 Passed: Memory retrieved with correct metadata.")

    print("\n--- All Simple Memory Tests Completed Successfully ---")
    #teardown_test_environment() # Commented out to keep the DB file for inspection

if __name__ == '__main__':
    run_simple_memory_test()
